use std::fs::File;
use std::io;
use std::io::{Read, Write};
use std::path::Path;
use std::time::Instant;

use ring::aead;
use ring::aead::Aad;
use ring::aead::Algorithm;
use ring::aead::BoundKey;
use ring::aead::Nonce;
use ring::aead::NonceSequence;
use ring::aead::OpeningKey;
use ring::aead::SealingKey;
use ring::aead::UnboundKey;
use ring::aead::AES_256_GCM;
use ring::aead::NONCE_LEN;
use ring::error::Unspecified;
use ring::rand::{SecureRandom, SystemRandom};

/// 对称加密算法AES使用演示

//AES_256_GCM加解密
//参考：https://web3developer.io/authenticated-encryption-in-rust-using-ring/
#[test]
fn test1() -> Result<(), Unspecified> {
    // 创建密钥

    let rand = SystemRandom::new();
    // 生成新的对称加密密钥
    let mut key_bytes = vec![0; AES_256_GCM.key_len()];
    rand.fill(&mut key_bytes)?;
    println!("key_bytes = {}", hex::encode(&key_bytes)); // don't print this in production code

    // 创建一个没有指定角色或随机序列的新 AEAD 密钥
    let unbound_key = UnboundKey::new(&AES_256_GCM, &key_bytes)?;

    // Create a new NonceSequence type which generates nonces
    let nonce_sequence = CounterNonceSequence::new(1);

    // Create a new AEAD key for encrypting and signing ("sealing"), bound to a nonce sequence
    // The SealingKey can be used multiple times, each time a new nonce will be used
    let mut sealing_key = SealingKey::new(unbound_key, nonce_sequence);

    //=======准备数据并使用 SealingKey 加密
    // This data will be authenticated but not encrypted
    //let associated_data = Aad::empty(); // is optional so can be empty
    let associated_data = Aad::from(b"additional public data");

    // 待加密数据
    let data = b"hello world";
    println!("data = {}", String::from_utf8(data.to_vec()).unwrap());

    // Create a mutable copy of the data that will be encrypted in place
    let mut in_out = data.clone();
    println!("in_out = {}", hex::encode(&in_out));

    // Encrypt the data with AEAD using the AES_256_GCM algorithm
    let tag = sealing_key.seal_in_place_separate_tag(associated_data, &mut in_out)?;
    println!("tag = {}", hex::encode(tag.as_ref()));
    println!("in_out = {}", hex::encode(&in_out));

    //====下面解密数据
    // 重新创建之前移动的变量
    let unbound_key = UnboundKey::new(&AES_256_GCM, &key_bytes)?;
    let nonce_sequence = CounterNonceSequence(1);
    //let associated_data = Aad::empty(); // supplying the wrong data causes the decryption to fail
    let associated_data = Aad::from(b"additional public data");

    // 创建新的 AEAD 密钥用于解密和验证身份验证标签
    let mut opening_key = OpeningKey::new(unbound_key, nonce_sequence);

    // 通过传入关联数据和附加身份验证标签的密文来解密数据
    let mut cypher_text_with_tag = [&in_out, tag.as_ref()].concat();
    let decrypted_data = opening_key.open_in_place(associated_data, &mut cypher_text_with_tag)?;
    println!(
        "decrypted_data = {}",
        String::from_utf8(decrypted_data.to_vec()).unwrap()
    );

    Ok(())
}

//生成对称加密密钥
//十六进制字符串
#[test]
fn test2() -> Result<(), Unspecified> {
    let rand = SystemRandom::new();
    // 生成新的对称加密密钥
    let mut key_bytes = vec![0; AES_256_GCM.key_len()];
    rand.fill(&mut key_bytes)?;
    println!("key_bytes = {}", hex::encode(&key_bytes)); // don't print this in production code

    Ok(())
}

//CHACHA20_POLY1305模式算法
#[test]
fn test3() {}

#[test]
fn test4() {
    let hex_key = "f0e5785cc29ad1597f52b889a8b4067f3877e981fb8c649bff01025a0a55ca4c";
    let associated_data = "c7dc65935ebc5041";

    let data = String::from("啦啦，啦啦啦，你好呀！");
    let encode_data = encode_str_aes_256_gcm(hex_key, associated_data, data.as_bytes()).unwrap();

    println!("加密前：{}\n加密后：{}", data, encode_data);

    let decode_data = decode_str_aes_256_gcm(hex_key, associated_data, &encode_data);
    match decode_data {
        Ok(d) => println!("解密后：{}", d),
        Err(e) => println!("解密失败：{}", e),
    }
}

//加解密文件
#[test]
fn test5() {
    let hex_key = "f0e5785cc29ad1597f52b889a8b4067f3877e981fb8c649bff01025a0a55ca4c";
    let associated_data = "c7dc65935ebc5041";

    //打开文件
    let mut file_data = vec![];
    File::open("../test.txt")
        .expect("打开文件失败")
        .read_to_end(&mut file_data)
        .expect("读取文件失败");

    let mut sealing_key: SealingKey<CounterNonceSequence> =
        make_key(&AES_256_GCM, &hex::decode(hex_key).unwrap(), 2);
    let tag = sealing_key
        .seal_in_place_separate_tag(Aad::from(associated_data.as_bytes()), &mut file_data)
        .unwrap();
    println!("tag = {}", hex::encode(tag.as_ref()));

    let encoded_file_data: Vec<u8> = [&file_data, tag.as_ref()].concat(); //最终加密数据
                                                                          //写出到新的加密文件
    let path = Path::new("test.byc");
    let mut file = File::create(path).expect("创建文件失败");
    match file.write_all(&encoded_file_data) {
        Ok(_) => println!("数据写入文件成功！[{}]", path.display()),
        Err(why) => println!("数据写入文件失败：{}", why),
    }

    //解密文件
    let mut encoded_file = File::open("test.byc").expect("打开加密文件失败");
    let mut encoded_file_data = vec![];
    encoded_file
        .read_to_end(&mut encoded_file_data)
        .expect("读取加密文件失败");

    let mut opening_key: OpeningKey<CounterNonceSequence> =
        make_key(&AES_256_GCM, &hex::decode(hex_key).unwrap(), 2);
    let decoded_file_data = opening_key
        .open_in_place(
            Aad::from(associated_data.as_bytes()),
            &mut encoded_file_data,
        )
        .unwrap();

    //写出到新的解密文件
    let path = Path::new("test_decode.txt");
    let mut file = File::create(path).expect("创建解密文件失败");
    match file.write_all(&decoded_file_data) {
        Ok(_) => println!("数据写入文件成功！[{}]", path.display()),
        Err(why) => println!("数据写入文件失败：{}", why),
    }

    /*if let Ok(encode_data) = encode_str_aes_256_gcm(hex_key, associated_data, &file_data) {
        println!("加密后：{}", encode_data);

        let decode_data = decode_str_aes_256_gcm(hex_key, associated_data, &encode_data);
        match decode_data {
            Ok(d) => {
                println!("解密后：{}", d);
            }
            _ => {}
        }
    }*/
}

#[test]
fn test6() {
    let hex_key = "f0e5785cc29ad1597f52b889a8b4067f3877e981fb8c649bff01025a0a55ca4c";
    let associated_data = "c7dc65935ebc5041";

    let start = Instant::now();
    encode_file_aes_256_gcm(
        hex_key,
        associated_data,
        "C:\\Users\\HG11\\.hgbio\\20231214084816_5e01324b8bcd58217e2ee8505419497c.zip",
        "C:\\Users\\HG11\\Desktop\\20231213231840_5e449ec6447e815a2a2fe8f61523ef2b-rs.byc",
    )
    .unwrap();
    let duration = start.elapsed();
    println!(">>>>加密文件持续时间：{:?}", duration);

    /*if let Ok(encode_data) = encode_file_aes_256_gcm(
        hex_key,
        associated_data,
        "test.txt",
        "C:\\Users\\HG11\\Desktop\\test111.byc",
    ) {
        println!("文件加密成功");
    }*/

    //=== 解密文件
    let start = Instant::now();
    decode_file_aes_256_gcm(
        hex_key,
        associated_data,
        "C:\\Users\\HG11\\Desktop\\20231212091228.byc",
        "C:\\Users\\HG11\\Desktop\\20231212091228.zip",
    )
    .unwrap();
    println!(">>>>解密文件持续时间：{:?}", start.elapsed());
}

/// 加密文件
fn encode_file_aes_256_gcm(
    hex_key: &str,
    associated_data: &str,
    ori_file_path: &str,
    encode_file_path: &str,
) -> io::Result<()> {
    let key_bytes = hex::decode(hex_key).expect("invalid hex key");
    let mut sealing_key: SealingKey<CounterNonceSequence> = make_key(&AES_256_GCM, &key_bytes, 1);
    let associated_data = Aad::from(associated_data.as_bytes());

    // 文件内容
    let path = Path::new(ori_file_path);
    let mut file = File::open(path).expect("打开待加密文件失败");
    let mut file_data = vec![];
    file.read_to_end(&mut file_data)
        .expect("读取待加密文件失败");

    //加密文件
    let tag = sealing_key
        .seal_in_place_separate_tag(associated_data, &mut file_data)
        .unwrap();
    let encoded_data = [&file_data, tag.as_ref()].concat();

    //把加密后的数据写出到文件
    let path = Path::new(encode_file_path);
    let mut file = File::create(path).expect("创建加密文件失败");
    file.write_all(&encoded_data).expect("写入加密文件失败");

    Ok(())
}

/// 解密文件
fn decode_file_aes_256_gcm(
    hex_key: &str,
    associated_data: &str,
    encode_file_path: &str,
    decode_file_path: &str,
) -> io::Result<()> {
    let key_bytes = hex::decode(hex_key).expect("invalid hex key");
    let associated_data = Aad::from(associated_data.as_bytes());

    //打开加密文件
    let path = Path::new(encode_file_path);
    let mut file = File::open(path).expect("打开加密文件失败");
    let mut encoded_file_data = vec![]; //存储加密的文件内容
    file.read_to_end(&mut encoded_file_data)
        .expect("读取加密文件失败");

    let mut opening_key: OpeningKey<CounterNonceSequence> = make_key(&AES_256_GCM, &key_bytes, 1);
    let decoded_file_data = opening_key
        .open_in_place(associated_data, &mut encoded_file_data)
        .unwrap();

    let path = Path::new(decode_file_path);
    let mut file = File::create(path).expect("创建解密文件失败");
    file.write_all(&decoded_file_data)
        .expect("写入解密文件失败");

    Ok(())
}

//AES_256_GCM加密字符串
fn encode_str_aes_256_gcm(
    hex_key: &str,
    associated_data: &str,
    data: &[u8],
) -> Result<String, Unspecified> {
    let key_bytes = hex::decode(hex_key).expect("invalid hex key");
    let mut sealing_key: SealingKey<CounterNonceSequence> = make_key(&AES_256_GCM, &key_bytes, 1);

    let associated_data = Aad::from(associated_data.as_bytes());

    let mut in_out = data.to_vec(); //创建一个可变副本

    let tag = sealing_key.seal_in_place_separate_tag(associated_data, &mut in_out)?;

    let encoded_data = [&in_out, tag.as_ref()].concat();

    Ok(hex::encode(encoded_data))
}

//AES_256_GCM解密字符串
fn decode_str_aes_256_gcm(
    hex_key: &str,
    associated_data: &str,
    encoded_data: &str,
) -> Result<String, Unspecified> {
    let key_bytes = hex::decode(hex_key).expect("invalid hex key");
    let associated_data = Aad::from(associated_data.as_bytes());

    let mut opening_key: OpeningKey<CounterNonceSequence> = make_key(&AES_256_GCM, &key_bytes, 1);

    let mut cypher_text_with_tag = hex::decode(encoded_data)
        .expect("invalid hex data")
        .to_vec();
    let decrypted_data = opening_key.open_in_place(associated_data, &mut cypher_text_with_tag)?;

    Ok(String::from_utf8(decrypted_data.to_vec()).unwrap())
}

fn make_key<K: BoundKey<CounterNonceSequence>>(
    algorithm: &'static Algorithm,
    key: &[u8],
    initial_value: u32,
) -> K {
    let key = UnboundKey::new(algorithm, key).unwrap();
    let nonce_sequence = CounterNonceSequence::new(initial_value);
    K::new(key, nonce_sequence)
}

fn make_less_safe_key(algorithm: &'static aead::Algorithm, key: &[u8]) -> aead::LessSafeKey {
    let key = aead::UnboundKey::new(algorithm, key).unwrap();
    aead::LessSafeKey::new(key)
}

/// 接下来，通过实现 trait 来创建 NonceSequence。在这个例子中，我使用一个简单的递增计数器作为生成随机数的方法。
/// 然后，我们创建一个新的 NonceSequence 实例，从 1 开始，并将它与 UnboundKey 一起传递，以创建一个新的 SealingKey。
/// SealingKey 是一种将密钥与随机序列绑定的类型。第一个用过的随机序列值将是 1，以此类推。
/// 请注意，创建新的 SealingKey 时，UnboundKey 和 NonceSequence 实例被移动。这种设计有助于防止重复使用这些敏感值
struct CounterNonceSequence(u32);

impl CounterNonceSequence {
    fn new(initial_value: u32) -> Self {
        Self(initial_value)
    }
}

impl NonceSequence for CounterNonceSequence {
    // called once for each seal operation
    fn advance(&mut self) -> Result<Nonce, Unspecified> {
        let mut nonce_bytes = vec![0; NONCE_LEN];

        let bytes = self.0.to_be_bytes();
        nonce_bytes[8..].copy_from_slice(&bytes);
        //println!("nonce_bytes = {}", hex::encode(&nonce_bytes));

        self.0 += 1; // advance the counter
        Nonce::try_assume_unique_for_key(&nonce_bytes)
    }
}
